import os
import base64
import requests
import multiprocessing as mp
import threading as td
import subprocess
from collections import Counter
import shutil

try:
    from cvio import cvio
except ImportError as e:
    print(e)
    from .cvio import cvio

def encodeImg(imgName):
    with open(imgName, 'rb') as f:
        data = f.read()
        encodestr = base64.b64encode(data)
    return str(encodestr, encoding='utf-8')

def GetUUID():
    cmd = 'wmic csproduct get uuid'
    uuid = str(subprocess.check_output(cmd))
    pos1 = uuid.find("\\n")+2
    uuid = uuid[pos1:-15]
    return uuid

def load_base_index(url, src, user, batch_size=128):
    url = '%s/load_index' % url
    images = cvio.load_image_list(src, recursive=True)
    n = len(images)
    m = 0
    j = 0
    for i in range(0, n, batch_size):
        batch_images, batch_labels = [], []
        batch = images[i:i+batch_size]
        m += len(batch)
        j += 1
        for img in batch:
            img = img.replace('/', os.sep).replace('\\', os.sep)
            batch_images.append(encodeImg(img))
            batch_labels.append(img.split(os.sep)[-2])
        rsp = requests.post(url, json={'imgs': batch_images, 'labels': batch_labels, 'user': user}).json()
        if rsp['status'] != 'ok':
            raise Exception(rsp['response'])
        else:
            print('[%d/%d] load batch %d ok' % (m, n, j))

def empty_cache(url, user):
    rsp = requests.post('%s/empty_cache' % url, json={'user': user}).json()
    if rsp['status'] != 'ok':
        raise Exception(rsp['response'])
    else:
        print('empty cache ok')

def auto_label_folder(url, base_dir, img_dir, save_dir, batch_size=128, score=0.5, topk=10, mode='move'):
    user = GetUUID()
    load_base_index(url, base_dir, user=user, batch_size=batch_size)
    images = cvio.load_image_list(img_dir)
    n = len(images)
    m = 0
    j = 0
    for i in range(0, n, batch_size):
        j += 1
        batch_images = images[i:i+batch_size]
        batch = [encodeImg(img) for img in batch_images]
        rsp = requests.post('%s/auto_label' % url, json={'imgs': batch, 'user': user, 'score': score, 'topk': topk}).json()
        if rsp['status'] != 'ok':
            raise Exception(rsp['response'])
        m += len(batch)
        print('[%d/%d] auto label batch %d ok' % (m, n, j))
        results = eval(rsp['response'])

        thread_list = []
        nc = len(batch_images)
        for i, img in enumerate(batch_images):
            # print(results)
            result = results[i]
            if not len(result):
                continue
            # vote
            label = Counter([l for i, l, s in result]).most_common(1)[0][0]
            # save
            save_path = os.path.join(save_dir, label)            
            t = td.Thread(target=custom_copy, args=(i+1, nc, img, save_path, mode))
            # t = mp.Process(target=custom_copy, args=(img, save_path))
            thread_list.append(t)
            
        num_workers = mp.cpu_count() * 4
        for i in range(0, len(thread_list), num_workers):
            threads = thread_list[i:i+num_workers]
            for t in threads:
                t.start()
            t.join()
    print()
    empty_cache(url, user)

def custom_copy(i, nc, src, dst, mode='move'):
    # print(src,flush=True, end='\r')
    print('[%d/%d] %s' % (i, nc, os.path.basename(src)))
    if not os.path.exists(dst):
        os.makedirs(dst)
    
    if mode == 'move':
        shutil.move(src, os.path.join(dst, os.path.basename(src)))
    else:
        shutil.copy(src, os.path.join(dst, os.path.basename(src)))

if __name__ == '__main__':
    url = 'http://192.168.11.25:9999'
    base_dir = r'C:\Users\Admin\Desktop\test\0729\cuts\yl_dp_dpty500P\base'
    img_dir = r'C:\Users\Admin\Desktop\test\0729\cuts\yl_dp_dpty500P\dest'
    batch_size = 128
    score = 0.1
    mode = 'copy'
    save_dir = r'C:\Users\Admin\Desktop\test\0729\cuts\yl_dp_dpty500P\save'
    auto_label_folder(url, base_dir, img_dir, save_dir, batch_size)